{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DhZ_xZrTKov5"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow.keras.layers import Input, Conv2D, GlobalAveragePooling2D, Dense\n",
        "from tensorflow.keras.models import Model\n",
        "from tensorflow.keras.activations import relu\n",
        "\n",
        "# Define an encoder using Convolutional layers\n",
        "def encoder(input_shape):\n",
        "    input_layer = Input(shape=input_shape)\n",
        "    conv1 = Conv2D(64, (3, 3), activation=relu, padding='same')(input_layer)\n",
        "    conv2 = Conv2D(128, (3, 3), activation=relu, padding='same')(conv1)\n",
        "    encoded = GlobalAveragePooling2D()(conv2)\n",
        "    return Model(inputs=input_layer, outputs=encoded, name='encoder')\n",
        "\n",
        "# Define a transformer model\n",
        "def transformer(input_shape):\n",
        "    input_layer = Input(shape=input_shape)\n",
        "    # Your transformer layers here\n",
        "    # Attention mechanisms, multi-head attention, positional encodings, etc.\n",
        "    transformed = # Transformer layers\n",
        "    return Model(inputs=input_layer, outputs=transformed, name='transformer')\n",
        "\n",
        "# Define a decoder using Dense layers\n",
        "def decoder(latent_dim, output_dim):\n",
        "    input_layer = Input(shape=(latent_dim,))\n",
        "    dense1 = Dense(128, activation=relu)(input_layer)\n",
        "    dense2 = Dense(256, activation=relu)(dense1)\n",
        "    output_layer = Dense(output_dim, activation='softmax')(dense2)\n",
        "    return Model(inputs=input_layer, outputs=output_layer, name='decoder')\n",
        "\n",
        "# Create an encoder-decoder-transformer model\n",
        "input_shape = (256, 256, 3)\n",
        "latent_dim = 512\n",
        "output_dim = 10  # Example for a classification task\n",
        "input_data = Input(shape=input_shape)\n",
        "encoded = encoder(input_shape)(input_data)\n",
        "transformed = transformer((latent_dim,))(encoded)\n",
        "decoded = decoder(latent_dim, output_dim)(transformed)\n",
        "\n",
        "model = Model(inputs=input_data, outputs=decoded, name='encoder_decoder_transformer')\n",
        "\n",
        "# Compile and train the model\n",
        "# You'll need to define your loss function, optimizer, and dataset for training\n",
        "\n",
        "# Example training loop:\n",
        "# model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])\n",
        "# model.fit(train_dataset, epochs=10, validation_data=val_dataset)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Cyclic Learning Rate Schedule class CyclicLR(Callback): def __init__(self, base_lr=1e-4, max_lr=3e-4, step_size=2000., mode='triangular2'): super(CyclicLR, self).__init__() self.base_lr = base_lr self.max_lr = max_lr self.step_size = step_size self.mode = mode self.clr_iterations = 0. self.trn_iterations = 0. if self.mode not in ['triangular', 'triangular2', 'exp_range']: raise KeyError(\"mode must be one of 'triangular', 'triangular2', or 'exp_range'\")"
      ],
      "metadata": {
        "id": "vTF-NinTLLlL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow.keras.layers import Input, Conv2DTranspose, Reshape, Dense\n",
        "from tensorflow.keras.models import Model\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import os\n",
        "import urllib.request\n",
        "import tarfile\n",
        "import shutil\n",
        "from vit_keras import vit\n",
        "from tensorflow.keras.callbacks import LearningRateScheduler, Callback, EarlyStopping\n",
        "from tensorflow.keras.regularizers import l2\n",
        "\n",
        "# Download STL-10 dataset (if not already downloaded)\n",
        "url = 'http://ai.stanford.edu/~acoates/stl10/stl10_binary.tar.gz'\n",
        "file_name = 'stl10_binary.tar.gz'\n",
        "\n",
        "if not os.path.exists(file_name):\n",
        "    urllib.request.urlretrieve(url, file_name)\n",
        "\n",
        "# Extract the dataset (if not already extracted)\n",
        "data_dir = 'stl10_binary'\n",
        "file_names = ['train_X.bin', 'train_y.bin', 'test_X.bin', 'test_y.bin']\n",
        "\n",
        "if os.path.exists(data_dir):\n",
        "    shutil.rmtree(data_dir)  # Delete the existing folder if it exists\n",
        "\n",
        "tar = tarfile.open(file_name, \"r:gz\")\n",
        "tar.extractall()\n",
        "tar.close()\n",
        "\n",
        "# Load the dataset\n",
        "x_train_path = os.path.join(data_dir, file_names[0])\n",
        "y_train_path = os.path.join(data_dir, file_names[1])\n",
        "x_test_path = os.path.join(data_dir, file_names[2])\n",
        "y_test_path = os.path.join(data_dir, file_names[3])\n",
        "\n",
        "x_train = np.fromfile(x_train_path, dtype=np.uint8).reshape(-1, 3, 96, 96).transpose(0, 2, 3, 1)\n",
        "y_train = np.fromfile(y_train_path, dtype=np.uint8) - 1\n",
        "x_test = np.fromfile(x_test_path, dtype=np.uint8).reshape(-1, 3, 96, 96).transpose(0, 2, 3, 1)\n",
        "y_test = np.fromfile(y_test_path, dtype=np.uint8) - 1\n",
        "\n",
        "# Preprocessing: normalize the data\n",
        "x_train = x_train.astype('float32') / 255.0\n",
        "x_test = x_test.astype('float32') / 255.0\n",
        "\n",
        "# Define the VIT Autoencoder architecture\n",
        "latent_dim = 128\n",
        "\n",
        "# Encoder (Vision Transformer)\n",
        "inputs = Input(shape=(96, 96, 3))\n",
        "x = vit.vit_l32(image_size=96, activation='gelu', pretrained=False, include_top=False, pretrained_top=False)(inputs)\n",
        "x = Reshape((-1, x.shape[-1]))(x)  # Flatten the sequence of patches\n",
        "x = tf.keras.layers.GlobalAveragePooling1D()(x)  # Reduce sequence to a single vector\n",
        "latent_space = Dense(latent_dim, activation='relu', kernel_regularizer=l2(1e-5), kernel_initializer='glorot_uniform')(x)  # Dense layer with L2 regularization\n",
        "\n",
        "# Decoder\n",
        "decoder_inputs = Input(shape=(latent_dim,))\n",
        "x = Dense(6 * 6 * 32, activation='relu', kernel_regularizer=l2(1e-5), kernel_initializer='glorot_uniform')(decoder_inputs)  # Dense layer with L2 regularization\n",
        "x = Reshape((6, 6, 32))(x)\n",
        "x = Conv2DTranspose(32, (3, 3), activation='relu', strides=(2, 2), padding='same', kernel_regularizer=l2(1e-5), kernel_initializer='glorot_uniform')(x)\n",
        "x = Conv2DTranspose(16, (3, 3), activation='relu', strides=(2, 2), padding='same', kernel_regularizer=l2(1e-5), kernel_initializer='glorot_uniform')(x)\n",
        "outputs = Conv2DTranspose(3, (3, 3), activation='sigmoid', padding='same', kernel_regularizer=l2(1e-5), kernel_initializer='glorot_uniform')(x)\n",
        "\n",
        "# AE model\n",
        "encoder = Model(inputs, latent_space)\n",
        "decoder = Model(decoder_inputs, outputs)\n",
        "\n",
        "# Create the autoencoder by connecting the encoder and decoder\n",
        "autoencoder_output = decoder(encoder(inputs))\n",
        "autoencoder = Model(inputs, autoencoder_output)\n",
        "\n",
        "# Reshape inputs and outputs\n",
        "inputs_reshaped = tf.image.resize(inputs, (24, 24))\n",
        "autoencoder_output_reshaped = tf.image.resize(autoencoder_output, (24, 24))\n",
        "\n",
        "# Define the loss function (Autoencoder loss)\n",
        "reconstruction_loss = tf.keras.losses.mean_squared_error(tf.reshape(inputs_reshaped, (-1, 24 * 24 * 3)),\n",
        "                                                         tf.reshape(autoencoder_output_reshaped, (-1, 24 * 24 * 3)))\n",
        "\n",
        "autoencoder_loss = tf.reduce_mean(reconstruction_loss)\n",
        "\n",
        "# Compile the model\n",
        "autoencoder.add_loss(autoencoder_loss)\n",
        "autoencoder.compile(optimizer='adam')\n",
        "\n",
        "# Cyclic Learning Rate Schedule\n",
        "class CyclicLR(Callback):\n",
        "    def __init__(self, base_lr=1e-4, max_lr=3e-4, step_size=2000., mode='triangular2'):\n",
        "        super(CyclicLR, self).__init__()\n",
        "        self.base_lr = base_lr\n",
        "        self.max_lr = max_lr\n",
        "        self.step_size = step_size\n",
        "        self.mode = mode\n",
        "        self.clr_iterations = 0.\n",
        "        self.trn_iterations = 0.\n",
        "        if self.mode not in ['triangular', 'triangular2', 'exp_range']:\n",
        "            raise KeyError(\"mode must be one of 'triangular', 'triangular2', or 'exp_range'\")\n",
        "\n",
        "    def clr(self):\n",
        "        cycle = np.floor(1 + self.clr_iterations / (2 * self.step_size))\n",
        "        x = np.abs(self.clr_iterations / self.step_size - 2 * cycle + 1)\n",
        "        if self.mode == 'triangular':\n",
        "            return self.base_lr + (self.max_lr - self.base_lr) * np.maximum(0, (1 - x))\n",
        "        elif self.mode == 'triangular2':\n",
        "            return self.base_lr + (self.max_lr - self.base_lr) * np.maximum(0, (1 - x)) / float(2  (cycle - 1))\n",
        "        elif self.mode == 'exp_range':\n",
        "            gamma = 1 / float(2  (cycle - 1))\n",
        "            return self.base_lr * (gamma  self.clr_iterations)\n",
        "\n",
        "    def on_train_begin(self, logs=None):\n",
        "        logs = logs or {}\n",
        "        tf.keras.backend.set_value(self.model.optimizer.lr, self.base_lr)\n",
        "\n",
        "    def on_batch_end(self, batch, logs=None):\n",
        "        logs = logs or {}\n",
        "        self.trn_iterations += 1\n",
        "        self.clr_iterations += 1\n",
        "        tf.keras.backend.set_value(self.model.optimizer.lr, self.clr())\n",
        "\n",
        "# Set up cyclic learning rate\n",
        "base_lr = 1e-5\n",
        "max_lr = 3e-4\n",
        "batch_size = 128\n",
        "step_size = 4 * (len(x_train) // batch_size)\n",
        "clr = CyclicLR(base_lr=base_lr, max_lr=max_lr, step_size=step_size, mode='triangular2')\n",
        "\n",
        "# Set up early stopping\n",
        "early_stopping = EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True)\n",
        "\n",
        "# Train the model\n",
        "epochs = 100\n",
        "batch_size = 128\n",
        "history = autoencoder.fit(x_train, epochs=epochs, batch_size=batch_size, validation_data=(x_test, None), callbacks=[clr, early_stopping])\n",
        "\n",
        "# Plot the learning curves for loss\n",
        "plt.figure(figsize=(10, 6))\n",
        "plt.plot(history.history['loss'], label='Train Loss')\n",
        "plt.plot(history.history['val_loss'], label='Validation Loss')\n",
        "plt.xlabel('Epochs')\n",
        "plt.ylabel('Loss')\n",
        "plt.title('Training and Validation Loss')\n",
        "plt.legend()\n",
        "plt.grid()\n",
        "plt.show()\n",
        "\n",
        "# Generate and plot some reconstructed images using the VIT Autoencoder\n",
        "num_samples = 5\n",
        "random_indices = np.random.randint(0, len(x_test), num_samples)\n",
        "sample_images = x_test[random_indices]\n",
        "reconstructed_images = autoencoder.predict(sample_images)\n",
        "\n",
        "plt.figure(figsize=(10, 4))\n",
        "\n",
        "for i in range(num_samples):\n",
        "    plt.subplot(2, num_samples, i + 1)\n",
        "    plt.imshow(sample_images[i])\n",
        "    plt.title(\"Original\")\n",
        "    plt.axis('off')\n",
        "\n",
        "    plt.subplot(2, num_samples, num_samples + i + 1)\n",
        "    plt.imshow(reconstructed_images[i])\n",
        "    plt.title(\"Reconstructed\")\n",
        "    plt.axis('off')\n",
        "\n",
        "plt.show()\n"
      ],
      "metadata": {
        "id": "c3no--eCLWcE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import timm  # Library for vision models, including SWIN\n",
        "\n",
        "# Create a simplified SWIN Transformer model\n",
        "class SwinTransformer(nn.Module):\n",
        "    def __init__(self, num_classes, img_size, patch_size, embed_dim, num_layers):\n",
        "        super(SwinTransformer, self).__init__()\n",
        "        self.patch_embed = timm.models.vision_transformer.patch_embed.PatchEmbed(\n",
        "            img_size=img_size, patch_size=patch_size, in_chans=3, embed_dim=embed_dim)\n",
        "\n",
        "        self.transformer = timm.models.vision_transformer.SwinTransformer(\n",
        "            img_size=img_size, embed_dim=embed_dim, depths=[num_layers] * 4, num_heads=embed_dim // 32,\n",
        "            num_classes=num_classes, qkv_bias=True)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.patch_embed(x)\n",
        "        x = self.transformer(x)\n",
        "        return x\n",
        "\n",
        "# Hyperparameters\n",
        "num_classes = 1000\n",
        "img_size = 224\n",
        "patch_size = 4\n",
        "embed_dim = 96\n",
        "num_layers = 12  # Number of layers in each stage of the SWIN Transformer\n",
        "\n",
        "# Create the SWIN Transformer model\n",
        "model = SwinTransformer(num_classes, img_size, patch_size, embed_dim, num_layers)\n",
        "\n",
        "# Input image tensor\n",
        "input_image = torch.randn(1, 3, img_size, img_size)  # Batch size 1, 3 channels (RGB), img_size x img_size\n",
        "\n",
        "# Forward pass through the model\n",
        "output = model(input_image)\n",
        "print(\"Output shape:\", output.shape)  # Print the shape of the output tensor\n"
      ],
      "metadata": {
        "id": "n744ZPtoLdZP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Develop VAE and AE  with SWIN Transformers\n",
        "\n",
        "# Import the required libraries\n",
        "import tensorflow as tf\n",
        "from tensorflow.keras.layers import Input, Conv2D, Flatten, Dense, Conv2DTranspose, Reshape, LayerNormalization\n",
        "from tensorflow.keras.models import Model\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "import tensorflow_addons as tfa\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import os\n",
        "import urllib.request\n",
        "from tensorflow.keras.regularizers import l2\n",
        "# Load the Swin Transformer model from Keras examples\n",
        "from keras.applications import SwinTransformer\n",
        "import urllib.request\n",
        "import tarfile\n",
        "from tensorflow.keras.callbacks import EarlyStopping\n",
        "\n",
        "# Download STL-10 dataset\n",
        "url = 'http://ai.stanford.edu/~acoates/stl10/stl10_binary.tar.gz'\n",
        "file_name = 'stl10_binary.tar.gz'\n",
        "\n",
        "if not os.path.exists(file_name):\n",
        "    urllib.request.urlretrieve(url, file_name)\n",
        "\n",
        "# Extract the dataset\n",
        "tar = tarfile.open(file_name, \"r:gz\")\n",
        "tar.extractall()\n",
        "tar.close()\n",
        "\n",
        "# Extract the dataset (if not already extracted)\n",
        "data_dir = 'stl10_binary'\n",
        "file_names = ['train_X.bin', 'train_y.bin', 'test_X.bin', 'test_y.bin']\n",
        "\n",
        "# Load the dataset\n",
        "x_train_path = os.path.join(data_dir, file_names[0])\n",
        "y_train_path = os.path.join(data_dir, file_names[1])\n",
        "x_test_path = os.path.join(data_dir, file_names[2])\n",
        "y_test_path = os.path.join(data_dir, file_names[3])\n",
        "\n",
        "x_train = np.fromfile(x_train_path, dtype=np.uint8).reshape(-1, 3, 96, 96).transpose(0, 2, 3, 1)\n",
        "y_train = np.fromfile(y_train_path, dtype=np.uint8) - 1\n",
        "x_test = np.fromfile(x_test_path, dtype=np.uint8).reshape(-1, 3, 96, 96).transpose(0, 2, 3, 1)\n",
        "y_test = np.fromfile(y_test_path, dtype=np.uint8) - 1\n",
        "\n",
        "# Initialize Swin Transformer model\n",
        "swin_model = SwinTransformer(\n",
        "    input_shape=(96, 96, 3),\n",
        "    num_classes=10,  # The number of classes doesn't matter for VAE and AE\n",
        "    num_heads=4,     # You can adjust the number of heads based on your requirement\n",
        "    num_transformer_blocks=2,  # You can adjust the number of blocks based on your requirement\n",
        ")\n",
        "\n",
        "# Remove the classification head from the model\n",
        "swin_model = tf.keras.Model(inputs=swin_model.inputs, outputs=swin_model.layers[-2].output)\n",
        "\n",
        "# Define the VAE architecture based on the Swin Transformer\n",
        "inputs = Input(shape=(96, 96, 3))\n",
        "x = swin_model(inputs)\n",
        "\n",
        "# Latent space\n",
        "latent_dim = 256\n",
        "z_mean = Dense(latent_dim, kernel_regularizer=l2(0.01))(x)  # Adding L2 regularization\n",
        "z_log_var = Dense(latent_dim, kernel_regularizer=l2(0.01))(x)  # Adding L2 regularization\n",
        "\n",
        "# Reparameterization trick\n",
        "def sampling(args):\n",
        "    z_mean, z_log_var = args\n",
        "    epsilon = tf.random.normal(shape=(tf.shape(z_mean)[0], latent_dim))\n",
        "    return z_mean + tf.exp(0.5 * z_log_var) * epsilon\n",
        "\n",
        "z = tf.keras.layers.Lambda(sampling)([z_mean, z_log_var])\n",
        "\n",
        "# Decoder\n",
        "decoder_inputs = Input(shape=(latent_dim,))\n",
        "x = Dense(6 * 6 * 64, activation='relu')(decoder_inputs)\n",
        "x = Reshape((6, 6, 64))(x)\n",
        "x = Conv2DTranspose(128, (3, 3), activation='relu', strides=(2, 2), padding='same')(x)\n",
        "x = Conv2DTranspose(64, (3, 3), activation='relu', strides=(2, 2), padding='same')(x)\n",
        "outputs = Conv2DTranspose(3, (3, 3), activation='sigmoid', padding='same')(x)\n",
        "\n",
        "# VAE model\n",
        "vae = Model(inputs, outputs)\n",
        "\n",
        "# Define the loss function for VAE\n",
        "def vae_loss(inputs, outputs):\n",
        "    reconstruction_loss = tf.keras.losses.binary_crossentropy(\n",
        "        tf.reshape(inputs, (-1, 24 * 24 * 3)), tf.reshape(outputs, (-1, 24 * 24 * 3))\n",
        "    )\n",
        "    reconstruction_loss *= 24 * 24 * 3\n",
        "    kl_loss = 1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)\n",
        "    kl_loss = tf.reduce_mean(kl_loss, axis=-1)\n",
        "    kl_loss *= -0.5\n",
        "    return tf.reduce_mean(reconstruction_loss + kl_loss)\n",
        "\n",
        "# Compile the VAE model\n",
        "vae.compile(optimizer=Adam(learning_rate=1e-4), loss=vae_loss)\n",
        "\n",
        "# AE model (without the latent space)\n",
        "encoder = Model(inputs, z_mean)\n",
        "decoder = Model(decoder_inputs, outputs)\n",
        "\n",
        "# Compile the AE model\n",
        "decoder.compile(optimizer=Adam(learning_rate=1e-4), loss='binary_crossentropy')\n",
        "\n",
        "\n",
        "# Early stopping\n",
        "early_stopping = EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True)\n",
        "\n",
        "# Training the VAE model\n",
        "epochs = 100\n",
        "batch_size = 128\n",
        "history_vae = vae.fit(x_train, x_train, epochs=epochs, batch_size=batch_size, validation_data=(x_test, x_test), callbacks=[early_stopping])\n",
        "\n",
        "# # Training the AE model\n",
        "# history_ae = decoder.fit(x_train, x_train, epochs=epochs, batch_size=batch_size, validation_data=(x_test, x_test), callbacks=[early_stopping])\n",
        "\n",
        "# # Save the models (optional)\n",
        "# vae.save(\"vae_model.h5\")\n",
        "# decoder.save(\"ae_model.h5\")\n",
        "\n",
        "# # Load the models (optional)\n",
        "# loaded_vae = tf.keras.models.load_model(\"vae_model.h5\")\n",
        "# loaded_decoder = tf.keras.models.load_model(\"ae_model.h5\")\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# Plot the learning curves for loss\n",
        "plt.figure(figsize=(10, 6))\n",
        "plt.plot(history_vae.history_vae['loss'], label='Train Loss')\n",
        "plt.plot(history_vae.history_vae['val_loss'], label='Validation Loss')\n",
        "plt.xlabel('Epochs')\n",
        "plt.ylabel('Loss')\n",
        "plt.title('Training and Validation Loss')\n",
        "plt.legend()\n",
        "plt.grid()\n",
        "plt.show()\n"
      ],
      "metadata": {
        "id": "z9eYCNXMLvF_"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}